#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/State.h>
#include <mavros_msgs/RCIn.h>
#include <nav_msgs/Path.h>
#include <string>
#include "math.h"
#include <offboard/ControlCommand.h>
#include <iostream>

#include "controller/colors.h"

/* -------------------------------------------------------------------------- */
/*                               Global Variable                              */
/* -------------------------------------------------------------------------- */
mavros_msgs::State cur_state;
geometry_msgs::PoseStamped cur_pose;
geometry_msgs::PoseStamped takeoff_pose;
geometry_msgs::PoseStamped hover_pose; // pose to hover in offboard mode
offboard::ControlCommand pva_traj;

double hover_height = 1.0;

int ctrl_rate = 25;

bool receive_traj = false;

/* -------------------------------------------------------------------------- */
/*                     Function Declaration and Definition                    */
/* -------------------------------------------------------------------------- */

// Tool Function
template <typename T>
void readParam(ros::NodeHandle &nh, std::string param_name, T &loaded_param)
{
  // template to read param from roslaunch
  const std::string &node_name = ros::this_node::getName();
  param_name = node_name + "/" + param_name;
  if (!nh.getParam(param_name, loaded_param))
  {
    std::cout << RED_B << "Fail to load " << param_name
              << ", use default value:" << loaded_param
              << RESET << std::endl;
  }
  else
  {
    std::cout << GREEN_B << "Load successfully " << param_name
              << ": " << loaded_param
              << RESET << std::endl;
  }
}

offboard::ControlCommand pose2pva(geometry_msgs::PoseStamped pose)
{
  offboard::ControlCommand pva;
  pva.header.stamp = ros::Time::now();
  pva.pos.x = pose.pose.position.x;
  pva.pos.y = pose.pose.position.y;
  pva.pos.z = pose.pose.position.z;
  pva.vel.x = 0.0;
  pva.vel.y = 0.0;
  pva.vel.z = 0.0;
  pva.acc.x = 0.0;
  pva.acc.y = 0.0;
  pva.acc.z = 0.0;
  pva.yaw = 0.0;
  return pva;
}

void loadRosParams(ros::NodeHandle &nh)
{
  readParam<int>(nh, "ctrl_rate", ctrl_rate);
  readParam<double>(nh, "hover_height", hover_height);
}

// callback function
void uav_state_cb(const mavros_msgs::State::ConstPtr &msg)
{
  cur_state = *msg;
}

void uav_pose_cb(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
  cur_pose = *msg;
}

void traj_cb(const offboard::ControlCommand::ConstPtr &msg)
{
  pva_traj = *msg;
  receive_traj = true;
}

class ReceiveRC
{
  struct Point_t
  {
    double x, y, z;
  };

public:
  ReceiveRC(ros::NodeHandle &nh)
  {
    rc_sub = nh.subscribe<mavros_msgs::RCIn>(
        "/mavros/rc/in", 10, &ReceiveRC::rcCallback, this);
    path_pub = nh.advertise<nav_msgs::Path>(
        "/drone_0_waypoint_generator/waypoints", 10);

    XmlRpc::XmlRpcValue points;
    const std::string &node_name = ros::this_node::getName();
    if (nh.getParam(node_name + "/" + "waypoints", points)) {
      ROS_ASSERT(points.getType() == XmlRpc::XmlRpcValue::TypeArray);

      for (int i = 0; i < points.size(); ++i) {
        ROS_ASSERT(points[i].getType() == XmlRpc::XmlRpcValue::TypeStruct);
        Point_t point;
        point.x = static_cast<double>(points[i]["x"]);
        point.y = static_cast<double>(points[i]["y"]);
        point.z = static_cast<double>(points[i]["z"]);
        waypoints.push_back(point);
      }
    }
    else {
      ROS_ERROR("Failed to get path points from the parameter server.");
    }

    for (const auto &point : waypoints) {
      ROS_INFO("Waypoint: %.2f, %.2f, %.2f", point.x, point.y, point.z);
    }
  }

  void rcCallback(const mavros_msgs::RCIn::ConstPtr &msg)
  {
    if (msg->channels[7] != 1934) return;

    ros::Time t_now = ros::Time::now();
    if ((t_now - last_request).toSec() < 5.0) {
      ROS_ERROR_THROTTLE(2.0, "Request too fast");
      return;
    }

    last_request = t_now;

    nav_msgs::Path path;
    path.header.stamp = t_now;
    geometry_msgs::PoseStamped pose;
    pose.pose.position.x = waypoints[index].x;
    pose.pose.position.y = waypoints[index].y;
    pose.pose.position.z = waypoints[index].z;
    path.poses.push_back(pose);
    path_pub.publish(path);
    index = (index + 1) % waypoints.size();
  }

private:
  ros::Subscriber rc_sub;
  ros::Publisher path_pub;
  ros::Time last_request{ros::Time(0)};
  std::vector<Point_t> waypoints;
  int index{0};
};

/* -------------------------------------------------------------------------- */
/*                                main function                               */
/* -------------------------------------------------------------------------- */

int main(int argc, char **argv)
{
  // ros init
  ros::init(argc, argv, "offboard_node", ros::init_options::AnonymousName);
  ros::NodeHandle nh;

  ReceiveRC reveiveRC(nh);

  // load param
  loadRosParams(nh);

  // ros pub and sub
  ros::Subscriber uav_state_sub = nh.subscribe<mavros_msgs::State>(
      "/mavros/state", 10, uav_state_cb);
  ros::Subscriber uav_pose_sub = nh.subscribe<geometry_msgs::PoseStamped>(
      "/mavros/vision_pose/pose", 2, uav_pose_cb);
  ros::Subscriber traj_sub = nh.subscribe<offboard::ControlCommand>(
      "/drone_0/control/pos_cmd", 10, traj_cb);
  ros::Publisher setpva_pub = nh.advertise<offboard::ControlCommand>(
      "/setpoint_pva", 10);

  // the setpoint publishing rate MUST be faster than 2Hz
  ros::Rate ctrl_loop(ctrl_rate);

  // wait for FCU connection
  while (ros::ok() && !cur_state.connected)
  {
    ros::spinOnce();
    ROS_INFO_THROTTLE(1, "wait for FCU connection");
    ctrl_loop.sleep();
  }

  /**
   * Pre Step :
   * Send a few setpoints before starting
   * Before entering Offboard mode,
   * you must have already started streaming setpoints.
   * Otherwise the mode switch will be rejected.
   */
  for (int i = 10; ros::ok() && i > 0; --i)
  {
    ros::spinOnce();
    setpva_pub.publish(pose2pva(cur_pose));
    ctrl_loop.sleep();
  }

  // takeoff: Record whether the uav has taken off
  // just_takeoff: check whether the uav change into offboard mode from other mode
  bool takeoff = false;
  bool just_takeoff = true;

  /**************************** main ctrl loop *******************************/
  while (ros::ok())
  {
    ctrl_loop.sleep();
    ros::spinOnce();

    if (!cur_state.armed) {
      ROS_INFO_STREAM_THROTTLE(1, " waiting for Vehicle arm");
      setpva_pub.publish(pose2pva(cur_pose));
      takeoff = false;
      just_takeoff = true;
      continue;
    }

    if (cur_state.mode != "OFFBOARD")
    {
      ROS_INFO_STREAM_THROTTLE(1, " waiting for Offboard cmd");
      setpva_pub.publish(pose2pva(cur_pose));
      just_takeoff = true;
      continue;
    }

    ROS_INFO_STREAM_THROTTLE(1, "\033[33m OFFBOARD Mode \033[0m");

    // Take off to specific height
    if (!takeoff)
    {
      takeoff_pose.pose.position = cur_pose.pose.position;
      takeoff_pose.pose.position.z = hover_height;

      geometry_msgs::PoseStamped climb_pose;
      climb_pose.pose.position = cur_pose.pose.position;

      double err = 0.1; // TODO: 高度阈值
      while (ros::ok() &&
             fabs(takeoff_pose.pose.position.z - cur_pose.pose.position.z) > err)
      {
        ros::spinOnce();
        // TODO:上升速度 0.2m/s
        climb_pose.pose.position.z = climb_pose.pose.position.z +
                                     0.2 / double(ctrl_rate);
        climb_pose.pose.position.z = std::min(climb_pose.pose.position.z,
                                              takeoff_pose.pose.position.z);

        setpva_pub.publish(pose2pva(climb_pose));
        ROS_INFO("[Takeoff]: Target: %.3f, %.3f, %.3f | Curr: %.3f, %.3f, %.3f",
                 climb_pose.pose.position.x,
                 climb_pose.pose.position.y,
                 climb_pose.pose.position.z,
                 cur_pose.pose.position.x,
                 cur_pose.pose.position.y,
                 cur_pose.pose.position.z);
        ctrl_loop.sleep();
      }
      ROS_INFO_STREAM("\033[33m Take off done \033[0m");
      takeoff = true;
      hover_pose = takeoff_pose;
    }

    // Switch into offboard mode just now, stay still for a while
    if (just_takeoff)
    {
      ROS_INFO_STREAM("\033[33m Hovering .. \033[0m");
      for (int i = ctrl_rate * 4; ros::ok() && i > 0; --i)
      {
        ros::spinOnce();
        setpva_pub.publish(pose2pva(takeoff_pose));
        ctrl_loop.sleep();
      }
      just_takeoff = false;
      hover_pose = takeoff_pose;
    }

    while (ros::ok() && !receive_traj)
    {
      ros::spinOnce();
      ROS_INFO_STREAM_THROTTLE(1, "waiting for trajectory");
      setpva_pub.publish(pose2pva(takeoff_pose));
      ctrl_loop.sleep();
    }

    ROS_INFO_STREAM_ONCE("\033[33m In Traj mode! \033[0m");
    setpva_pub.publish(pva_traj);
  }

  return 0;
}
